#include "ConnectContext.h"
#include <Mswsock.h>
#include "Proxy.h"
#include "debug.h"


//////////////////////////////////////////////////////////////////////////
//  class BridgeContext
//
BridgeContext::BridgeContext(ConnectContext& connectContext, SOCKET& inSocket, SOCKET&  outSocket)
    : WSAOVERLAPPED(), connectContext(connectContext), complitedOperation(IOAccept),
      _inSocket(inSocket), _outSocket(outSocket)
{
}

int BridgeContext::recv()
{
    complitedOperation = IORead;

    WSABUF recvBuff;
    recvBuff.buf = buff;
    recvBuff.len = sizeof(buff);

    DWORD flags = 0;
    DWORD bytesRecv;
    return WSARecv(_inSocket, &recvBuff, 1, &bytesRecv, &flags, this, NULL);
}

#ifdef _DEBUG

#include <fstream>
using namespace std;

void BridgeContext::_dbgPrintBuff(DWORD bytesRecv)
{
    char name[MAX_PATH];
    sprintf(name, "logs/%i-%i.log", (int)_inSocket, (int)_outSocket);
    ofstream os(name, ios::out|ios::app);
    if(os.is_open()) {
        os.write(buff, bytesRecv);
        os << endl << "----------------------------------" << endl;
    }
    else {
        dbgprint("Can't open %s", name);
    }
}

#else

void BridgeContext::_dbgPrintBuff(DWORD bytesRecv) {}

#endif //_DEBUG


int BridgeContext::send(DWORD bytesRecv)
{
    _dbgPrintBuff(bytesRecv);

    complitedOperation = IOWrite;

    WSABUF sendBuff;
    sendBuff.buf = buff;
    sendBuff.len = bytesRecv;

    DWORD bytesSend = 0;
    DWORD flags = 0;
    return WSASend(_outSocket, &sendBuff, 1, &bytesSend, flags, this, NULL);
}

void BridgeContext::close()
{
    int ret = shutdown(_inSocket, SD_RECEIVE);
    if(ret == SOCKET_ERROR ) {
        dbgprint("Error occurred while shutdown SD_RECEIVE: %d.", WSAGetLastError());
    }
    ret = shutdown(_outSocket, SD_SEND);
    if(ret == SOCKET_ERROR ) {
        dbgprint("Error occurred while shutdown SD_SEND: %d.", WSAGetLastError());
    }
}


//////////////////////////////////////////////////////////////////////////
//  class ConnectContext
//
ConnectContext::ConnectContext() 
    : _userSocket(INVALID_SOCKET), _serverSocket(INVALID_SOCKET),
      _userToServerBridge(*this, _userSocket, _serverSocket),
      _serverToUserBridge(*this, _serverSocket, _userSocket),
      _activeBridgeCount(2)
{

}

ConnectContext::~ConnectContext()
{
    _close();
}

bool ConnectContext::initialize(HANDLE iocp)
{
    _userSocket = WSASocket(AF_INET, SOCK_STREAM, 0/*IPPROTO_IP*/, NULL, 0, WSA_FLAG_OVERLAPPED);
    if(_userSocket == INVALID_SOCKET) {
        dbgprint("Error occurred while create userSocket: %d.", WSAGetLastError());
        return false;
    }

    _serverSocket = WSASocket(AF_INET, SOCK_STREAM, 0/*IPPROTO_IP*/, NULL, 0, WSA_FLAG_OVERLAPPED);
    if(_serverSocket == INVALID_SOCKET) {
        dbgprint("Error occurred while create serverSocket: %d.", WSAGetLastError());
        return false;
    }

    HANDLE temp = CreateIoCompletionPort((HANDLE)_userSocket, iocp, (DWORD)this, 0);
    if(temp == NULL) {
        dbgprint("Error occurred while executing CreateIoCompletionPort(): %d", GetLastError());
        return false;
    }
    temp = CreateIoCompletionPort((HANDLE)_serverSocket, iocp, (DWORD)this, 0);
    if(temp == NULL) {
        dbgprint("Error occurred while executing CreateIoCompletionPort(): %d", GetLastError());
        return false;
    }

    return true;
}

bool ConnectContext::close(BridgeContext* bridge)
{
    if(_activeBridgeCount == 0) {
        return true;
    }

    bridge->close();
    _activeBridgeCount--;
    if(_activeBridgeCount == 0) {
        _close();
        return true;
    }

    return false;
}

void ConnectContext::_close()
{
    if(_userSocket != INVALID_SOCKET) {
        int ret = closesocket(_userSocket);
        if(ret == SOCKET_ERROR) {
            dbgprint("Error occurred while close _userSocket: %d.", WSAGetLastError());
        }
        else {
            _userSocket = INVALID_SOCKET;
        }
    }
    if(_serverSocket != INVALID_SOCKET) {
        int ret = closesocket(_serverSocket);
        if(ret == SOCKET_ERROR ) {
            dbgprint("Error occurred while close _serverSocket: %d.", WSAGetLastError());
        }
        else {
            _serverSocket = INVALID_SOCKET;
        }
    }
}

bool ConnectContext::waitAccept(SOCKET listenSocket)
{
    DWORD bytes;
    BOOL res = AcceptEx(listenSocket, _userSocket, _userToServerBridge.getBuff(),
        _userToServerBridge.getBuffSize() - 2 * (sizeof(sockaddr_in) + 16),
        sizeof(sockaddr_in) + 16, sizeof(sockaddr_in) + 16,
        &bytes, &_userToServerBridge);
    if(res == FALSE && WSAGetLastError() != WSA_IO_PENDING) {
        dbgprint("Error occurred while executing AcceptEx(): %d", WSAGetLastError());
        return false;
    }

    return true;
}

#include <string>
using namespace std;
void getTarget(const char* request, string& host, string& port);

int ConnectContext::completeAccept(SOCKET listenSocket, DWORD bytesRecv)
{
    setsockopt(_userSocket, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT, (char*)&listenSocket, sizeof(listenSocket));

#if 0
    sockaddr_in address = {0};
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = inet_addr("213.180.204.8");   // !!! ya.ru
    address.sin_port = htons(80);    
#else    
    string host, port;
    getTarget(_userToServerBridge.getBuff(), host, port);
    hostent* pHost = gethostbyname(host.c_str());

    sockaddr_in address = {0};
    address.sin_family = AF_INET;
    memcpy(&address.sin_addr.s_addr, pHost->h_addr, pHost->h_length);
    address.sin_port = htons( atoi(port.c_str()) );
#endif

    int ret = connect(_serverSocket, (sockaddr*)&address, sizeof(address));
    if ((SOCKET_ERROR == ret) && (WSA_IO_PENDING != WSAGetLastError())) {
        dbgprint("Error occurred while connecting: %d.", WSAGetLastError());
        return ret;
    }

    ret = _serverToUserBridge.recv();
    if ((SOCKET_ERROR == ret) && (WSA_IO_PENDING != WSAGetLastError())) {
        dbgprint("Error occurred in _serverToUserBridge.recv(): %d.", WSAGetLastError());
        return ret;
    }

    return _userToServerBridge.send(bytesRecv);
}

void getTarget(const char* request, string& host, string& port)
{
    host="";
    port="";

    //skip the first HTTP command
    while(*request!=' ')
    {
        ++request;
    }
    //skip space
    while( (*request<'a' || *request>'z') && (*request<'A' || *request>'Z') )
        ++request;
    const char* urlStart=request;

    if(    (*request=='H' || *request=='h')
        && (*request=='T' || *request=='t')
        && (*request=='T' || *request=='t')
        && (*request=='P' || *request=='p')
        && (*request==':')
        )
    {
        while(' '==*request)
            ++request;

        while(' '!=*request && ':'!=*request && '\0'!=*request && '\r'!=*request && '\n'!=*request)
        {
            host+=*request;
            ++request;
        }
        if(':'==*request)
        {
            while(' '!=*request && ':'!=*request && '\0'!=*request && '\r'!=*request && '\n'!=*request)
            {
                port+=*request;
                ++request;
            }
        }
        if(port.length()==0)
            port="80";

        return;
    }

    while(*request!='\0')
    {
        //Process "Host: "
        if(*request=='\n' && *(request+1)=='H' && *(request+2)=='o' && *(request+3)=='s' && *(request+4)=='t' && *(request+5)==':')
        {
            request+=6;

            while( (*request<'a' || *request>'z') && (*request<'A' || *request>'Z') )
                ++request;
            //find host address in "Host: "
            while( *request!=' ' && *request!='\t' && *request!='\r' && *request!='\n' && *request!='\0' && *request!=':')
            {
                host+=*request;
                ++request;
            }
            //find port number in "Host: "
            if(*request==':')
            {
                while( *request!=' ' && *request!='\t' && *request!='\r' && *request!='\n' && *request!='\0')
                {
                    port+=*request;
                    ++request;
                }
            }
            break;
        }
        ++request;
    }

    if(host.length()>0 && port.length()==0)
    {
        while(' '!=*urlStart)
        {
            if(':'==*urlStart)
            {
                ++urlStart;
                while(' '!=*urlStart && '/'!=*urlStart && '\\'!=*urlStart)
                    port+=*urlStart++;
                break;
            }
            ++urlStart;
        }

        if(port.length()==0)
            port="80";

        return;
    }
}
